__author__ = 'chen_ju0606@163.com'

import os.path
import sys
import time
from datetime import datetime

import keyboard
from PIL import ImageGrab
from PyQt5.QtGui import QGuiApplication, QPixmap, QPainter, QPen, QColor, QMouseEvent, QKeySequence
from PyQt5.QtWidgets import QWidget, QVBoxLayout, QPushButton, QApplication, QMainWindow, QRubberBand, QLabel, \
    QToolBar, QComboBox, QMessageBox, QShortcut, QHBoxLayout
from PyQt5.QtCore import Qt, QPoint, QRect, QSize

from tools.files_helper import extract_case_code_list, classify_images_by_title, create_docx_with_titles_and_images, \
    back_up_images
from config.readconfig import ini
from tools.logger import log


class ScreenshotTool(QWidget):
    def __init__(self, start_case_code=None, shot_index=1):
        super().__init__()
        self.case_code_list = extract_case_code_list('input')
        self.start_case_code = start_case_code
        self.shot_index = shot_index
        self.initUI()

        if len(self.case_code_list) == 0:
            log.error(f'请检查输入文件input.docx或input.xlsx是否存在或文件内容是否错误')
            self.show_error_message(f'请检查文件input.docx或input.xlsx无误后重启应用程序')

        log.info(f'工具初始化完成!')

    def initUI(self):
        self.setWindowTitle('Screenshot Tool')

        # 置顶显示窗口
        self.setWindowFlags(self.windowFlags() | Qt.WindowStaysOnTopHint)

        layout = QVBoxLayout()

        base_tool_row = QHBoxLayout()
        shot_tool_row = QHBoxLayout()

        # Create a combo box
        self.combo_box = QComboBox(self)
        self.combo_box.setEditable(True)
        editor = self.combo_box.lineEdit()
        editor.setPlaceholderText('Enter Or Select Start Case Code')
        self.combo_box.addItems(self.case_code_list)
        base_tool_row.addWidget(self.combo_box)

        # Button for Start
        self.btn_start = QPushButton('Start', self)
        self.btn_start.clicked.connect(self.start)
        base_tool_row.addWidget(self.btn_start)

        # Button for reporting
        btn_report = QPushButton('Report', self)
        btn_report.clicked.connect(self.report)
        base_tool_row.addWidget(btn_report)

        # Button for full screen screenshot
        self.btn_fullscreen = QPushButton('Full Screen', self)
        self.btn_fullscreen.clicked.connect(self.fullscreen_screenshot)
        self.btn_fullscreen.setEnabled(False)
        shot_tool_row.addWidget(self.btn_fullscreen)

        # Button for custom area screenshot
        self.btn_custom = QPushButton('Custom Area', self)
        self.btn_custom.clicked.connect(self.custom_area_screenshot)
        self.btn_custom.setEnabled(False)
        shot_tool_row.addWidget(self.btn_custom)

        # Button for finishing and saving
        self.btn_finish = QPushButton('Finish', self)
        self.btn_finish.clicked.connect(self.finish)
        self.btn_finish.setEnabled(False)
        shot_tool_row.addWidget(self.btn_finish)

        layout.addLayout(base_tool_row)
        layout.addLayout(shot_tool_row)

        self.setLayout(layout)
        self.show()

    def mousePressEvent(self, event: QMouseEvent):
        if event.button() == Qt.LeftButton and event.pos().y() < 40:  # Assuming title bar height is 40
            self.drag_start_position = event.globalPos() - self.pos()
            self.dragging = True
            event.accept()

    def mouseMoveEvent(self, event: QMouseEvent):
        if self.dragging:
            self.move(event.globalPos() - self.drag_start_position)
            event.accept()

    def mouseReleaseEvent(self, event: QMouseEvent):
        self.dragging = False

    def start(self):
        self.start_case_code = self.combo_box.currentText()
        log.info(f'Start Case Code成功: {self.start_case_code}')

        if self.start_case_code not in self.case_code_list:
            log.info(f'数据错误: 列表中不存在{self.start_case_code}')
            self.show_error_message(f'数据错误: 列表中不存在{self.start_case_code}')
            # raise ValueError(f'数据错误: 列表中不存在{self.start_case_code}')
            return
        self.btn_start.setCheckable(False)
        self.btn_fullscreen.setEnabled(True)
        self.btn_custom.setEnabled(True)
        self.btn_finish.setEnabled(True)
        self.show_info_message(f'即将从案例编号: {self.start_case_code} 开始执行')

    def fullscreen_screenshot(self):
        if self.start_case_code == '' or self.start_case_code is None or not self.btn_fullscreen.isEnabled():
            # self.show_error_message("请选择或输入Start Case Code并点击Start按钮")
            return
        self.hide()
        time.sleep(0.2)
        screenshot = ImageGrab.grab()
        if not os.path.exists('images'):
            os.makedirs('images')
        screenshot_path = f'./images/screenshot_{self.start_case_code}_{str(self.shot_index).zfill(3)}.png'
        screenshot.save(screenshot_path)
        log.info(f'完成全屏截图, 已存放到images/screenshot_{self.start_case_code}_{str(self.shot_index).zfill(3)}.png')
        self.shot_index = self.shot_index + 1
        self.show()

    def custom_area_screenshot(self):
        if self.start_case_code == '' or self.start_case_code is None or not self.btn_custom.isEnabled():
            # self.show_error_message("请选择或输入Start Case Code并点击Start按钮")
            return
        self.hide()  # Hide the window while selecting the area
        self.selection_window = SelectionWindow(self, self.start_case_code, self.shot_index)
        self.selection_window.show()
        log.info(f'打开自定义区域选择')
        self.shot_index = self.shot_index + 1

    def finish(self):
        try:
            if self.start_case_code == '' or self.start_case_code is None or not self.btn_finish.isEnabled():
                # self.show_error_message("请选择或输入Start Case Code并点击Start按钮")
                return
            # 当到达最后一个元素时
            if self.start_case_code == self.case_code_list[len(self.case_code_list) - 1]:
                log.info(f'已完成所有测试用例')
                self.show_info_message(f'已完成所有测试用例')
                return

            log.info(f'当前案例{self.start_case_code}执行完成')
            self.start_case_code = self.case_code_list[self.case_code_list.index(self.start_case_code) + 1]
            self.shot_index = 1

            if self.start_case_code is not None:
                # self.input_box.setText(self.start_case_code)
                index = self.combo_box.findText(self.start_case_code)
                if index != -1:
                    self.combo_box.setCurrentIndex(index)

            self.btn_start.setCheckable(True)
            log.info(f'开始执行案例{self.start_case_code}')
        except Exception as e:
            print(f'错误信息: {e}')

    def report(self):
        # 将图片按照标题分类
        titles_images = classify_images_by_title('images', self.case_code_list)

        # 创建新的 .docx 文件
        docx_file_name = f"traces_out_{datetime.now().strftime('%Y%m%d%H%M%S%f')}.docx"
        create_docx_with_titles_and_images(titles_images, docx_file_name)
        # 备份照片文件
        tar_dir = f"./images_backup/{datetime.now().strftime('%Y%m%d%H%M%S%f')}"
        if not int(ini.is_backup_images) == 0:
            back_up_images('images', tar_dir)
            log.info(f'已完成图片备份，备份地址: {tar_dir}')
        self.show_info_message(f'已生成证迹文件: {docx_file_name}')

    def show_error_message(self, message):
        # Create a QMessageBox for error messages
        msg_box = QMessageBox(self)
        msg_box.setIcon(QMessageBox.Critical)
        msg_box.setText(message)
        msg_box.setWindowTitle('Error')
        msg_box.setStandardButtons(QMessageBox.Ok)
        msg_box.exec_()

    def show_info_message(self, message):
        # Create a QMessageBox for error messages
        msg_box = QMessageBox(self)
        msg_box.setIcon(QMessageBox.Information)
        msg_box.setText(message)
        msg_box.setWindowTitle('Info')
        msg_box.setStandardButtons(QMessageBox.Ok)


        msg_box.setWindowFlags(msg_box.windowFlags() | Qt.WindowStaysOnTopHint)
        msg_box.exec_()


class SelectionWindow(QMainWindow):
    def __init__(self, parent=None, start_case_code=None, index=None):
        super().__init__(parent)
        self.parent = parent
        self.start_case_code = start_case_code
        self.index = index
        self.setWindowTitle('Select Area')
        self.setWindowState(Qt.WindowFullScreen)
        self.setWindowOpacity(0.3)  # Make the selection window semi-transparent
        self.rubberBand = QRubberBand(QRubberBand.Rectangle, self)
        self.origin = QPoint()
        self.label = QLabel(self)
        self.label.setGeometry(10, 10, 300, 30)
        self.label.setText('Click and drag to select area, then release.')

        # 设置Esc键退出
        self.shortcut_esc = QShortcut(QKeySequence('Esc'), self)
        self.shortcut_esc.activated.connect(self.close_selection)

    def close_selection(self):
        self.close()
        self.parent.show()
        log.info(f'取消截图并返回主窗口')

    def mousePressEvent(self, event):
        self.origin = event.pos()
        self.rubberBand.setGeometry(QRect(self.origin, QSize()))
        self.rubberBand.show()

    def mouseMoveEvent(self, event):
        self.rubberBand.setGeometry(QRect(self.origin, event.pos()).normalized())

    def mouseReleaseEvent(self, event):
        rect = self.rubberBand.geometry()
        self.rubberBand.hide()
        screenshot = QGuiApplication.primaryScreen().grabWindow(0, rect.x(), rect.y(), rect.width(), rect.height())

        # Convert the screenshot to a QPixmap for editing
        pixmap = QPixmap(screenshot)

        self.close()
        log.info(f'关闭区域选择遮罩层')

        # Open an editor window with the screenshot
        self.editor_window = ImageEditorWindow(pixmap, self.parent, self.start_case_code, self.index)
        self.editor_window.show()
        log.info(f'展示图片编辑窗口')


class ImageEditorWindow(QMainWindow):
    def __init__(self, pixmap, parent=None, start_case_code=None, shot_index=None):
        super().__init__(parent)
        self.parent = parent
        self.start_case_code = start_case_code
        self.shot_index = shot_index
        self.setWindowTitle('Edit Screenshot')

        # Main layout for the window
        main_layout = QVBoxLayout()

        # Create a toolbar at the top
        self.toolbar = QToolBar('Edit', self)
        self.addToolBar(Qt.TopToolBarArea, self.toolbar)  # 将工具栏放置在顶部

        # Add the rectangle button to the toolbar
        self.rect_button = QPushButton('Draw Rectangle', self)
        self.rect_button.setCheckable(True)  # 设置为可选中状态
        self.rect_button.clicked.connect(self.toggle_rectangle_mode)
        self.toolbar.addWidget(self.rect_button)

        # Add the finish button to the toolbar
        self.finish_button = QPushButton('Finish', self)
        self.finish_button.clicked.connect(self.finish_editing)
        self.toolbar.addWidget(self.finish_button)

        # Create the label to hold the image
        self.label = QLabel(self)
        self.label.setPixmap(pixmap)

        # Add the image label to the layout
        main_layout.addWidget(self.label)

        # Set the central widget
        central_widget = QWidget()
        central_widget.setLayout(main_layout)
        self.setCentralWidget(central_widget)

        # Adjust window size based on contents
        self.adjustSize()

        self.drawing_rect = False
        self.rect_origin = QPoint()
        self.current_rect = QRect()

    def closeEvent(self, event):
        self.parent.show()  # 在关闭编辑窗口时显示主窗口
        log.info(f'关闭编辑窗口并返回主窗口')
        event.accept()  # 允许窗口关闭

    def toggle_rectangle_mode(self):
        # Toggle the drawing rectangle mode
        if self.rect_button.isChecked():
            self.drawing_rect = True
        else:
            self.drawing_rect = False

    def mousePressEvent(self, event):
        if self.drawing_rect and event.button() == Qt.LeftButton:
            self.rect_origin = event.pos() - self.label.pos() - QPoint(0, self.toolbar.height())  # 调整起始点，减去工具栏高度
            self.current_rect = QRect(self.rect_origin, QSize())

    def mouseMoveEvent(self, event):
        if self.drawing_rect and event.buttons() == Qt.LeftButton:
            self.current_rect.setBottomRight(
                event.pos() - self.label.pos() - QPoint(0, self.toolbar.height()))  # 调整位置，减去工具栏高度
            self.update()  # 触发重绘

    def mouseReleaseEvent(self, event):
        if self.drawing_rect and event.button() == Qt.LeftButton:
            self.current_rect.setBottomRight(
                event.pos() - self.label.pos() - QPoint(0, self.toolbar.height()))  # 调整位置，减去工具栏高度

            # 绘制最终的矩形
            painter = QPainter(self.label.pixmap())
            pen = QPen(QColor('red'), 3, Qt.SolidLine)
            painter.setPen(pen)
            painter.drawRect(self.current_rect)
            self.label.update()

    def finish_editing(self):
        if not os.path.exists('images'):
            os.makedirs('images')

        screenshot_path = f'./images/screenshot_{self.start_case_code}_{str(self.shot_index).zfill(3)}.png'
        self.label.pixmap().save(screenshot_path)

        log.info(
            f'完成自定义区域截图, 已存放到images/screenshot_{self.start_case_code}_{str(self.shot_index).zfill(3)}.png')

        self.close()
        log.info(f'关闭文件编辑窗口')

        self.parent.show()
        log.info(f'打开主窗口')


def main():
    try:
        app = QApplication(sys.argv)
        tool = ScreenshotTool()

        # Set global hotkeys
        keyboard.add_hotkey(ini.full_screenshot, tool.fullscreen_screenshot)  # Example for full screen screenshot
        # keyboard.add_hotkey(ini.custom_screenshot, tool.custom_area_screenshot)  # Example for full screen screenshot
        keyboard.add_hotkey(ini.finish, tool.finish)  # Example for finish

        sys.exit(app.exec_())
    except Exception as e:
        print(f"An error occurred: {e}")
        sys.exit(1)  # Terminate the process with a non-zero exit code


if __name__ == '__main__':
    main()
